"""
均值回归债券交易策略
当债券价格低于过去100天的均值-2个标准差时买入
当债券价格高于过去100天的均值+2个标准差时卖出
"""
from datetime import datetime
import numpy as np
import pandas as pd
from loguru import logger

from engine.event import (
    EventType, OrderEvent, OrderType, AssetType, 
    CashFlowEvent, CashFlowType
)
from strategy.base import Strategy

class BondMeanReversionStrategy(Strategy):
    """
    均值回归债券交易策略
    
    参数:
        name: 策略名称
        lookback_period: 回看天数，默认100天
        std_multiplier: 标准差倍数，默认2
        leverage: 杠杆倍数，默认1
    """
    
    def __init__(
        self, 
        name: str = "债券均值回归策略",
        lookback_period: int = 100,
        std_multiplier: float = 2.0,
        leverage: float = 1.0
    ):
        super().__init__(name=name)
        self.lookback_period = lookback_period
        self.std_multiplier = std_multiplier
        self.leverage = leverage
        self.price_history = {}  # 用于存储每个资产的价格历史
        logger.info(
            f"初始化策略: {self.name}, "
            f"回看期={lookback_period}, "
            f"标准差倍数={std_multiplier}, "
            f"杠杆={leverage}"
        )
    
    def initialize(self) -> None:
        """初始化策略"""
        # 初始化价格历史字典
        for symbol in self.engine.market_data.keys():
            if self.engine.assets[symbol].asset_type == AssetType.BOND:
                self.price_history[symbol] = []
    
    def update_price_history(self, symbol: str, price: float) -> None:
        """
        更新价格历史
        
        Args:
            symbol: 资产代码
            price: 当前价格
        """
        self.price_history[symbol].append(price)
        # 只保留lookback_period天的数据
        if len(self.price_history[symbol]) > self.lookback_period:
            self.price_history[symbol].pop(0)
    
    def calculate_signals(self, symbol: str, current_price: float) -> int:
        """
        计算交易信号
        
        Args:
            symbol: 资产代码
            current_price: 当前价格
            
        Returns:
            1: 买入信号
            -1: 卖出信号
            0: 无信号
        """
        if len(self.price_history[symbol]) < self.lookback_period:
            return 0
        
        # 计算均值和标准差
        price_array = np.array(self.price_history[symbol])
        mean_price = np.mean(price_array)
        std_price = np.std(price_array)
        
        # 计算上下轨
        upper_band = mean_price + self.std_multiplier * std_price
        lower_band = mean_price - self.std_multiplier * std_price
        
        # 生成信号
        if current_price < lower_band:
            return 1  # 买入信号
        elif current_price > upper_band:
            return -1  # 卖出信号
        return 0  # 无信号
    
    def on_data(self, market_data: dict) -> None:
        """
        处理市场数据
        
        Args:
            market_data: 市场数据字典
        """
        current_time = self.engine.current_time
        available_cash = self.engine.position.current_cash
        
        # 遍历所有资产
        for symbol, data in market_data.items():
            # 只处理债券资产
            asset = self.engine.assets.get(symbol)
            if asset.asset_type != AssetType.BOND:
                continue
            
            # 获取当前价格并更新价格历史,债券全价
            
            current_price = data.loc[current_time, 'close'] 
            current_price_full = current_price + asset.accrued_interest(current_time, 1)
            self.update_price_history(symbol, current_price)
            
            # 获取当前持仓
            position = self.engine.position.positions.get(symbol)
            current_quantity = position.quantity if position else 0
            
            # 计算交易信号
            signal = self.calculate_signals(symbol, current_price)
            
            if signal != 0:  # 有交易信号
                if signal > 0 and current_quantity == 0:  # 买入信号且当前无持仓
                    # 计算可买数量
                    max_position = available_cash * self.leverage
                    quantity = int(max_position / current_price_full)
                    
                    if quantity > 0:
                        # 如果使用杠杆，先进行回购融资
                        if self.leverage > 1:
                            repo_amount = available_cash * (self.leverage - 1)
                            for repo_symbol, repo_data in market_data.items():
                                repo_asset = self.engine.assets.get(repo_symbol)
                                if repo_asset.asset_type != AssetType.REPO:
                                    continue
                                
                                # 执行回购融资
                                repo_order = OrderData(
                                    type=EventType.ORDER,
                                    timestamp=current_time,
                                    asset_type=AssetType.REPO,
                                    symbol=repo_symbol,
                                    quantity=repo_amount,
                                    direction=-1,  # -1表示借入
                                    order_type=OrderType.MARKET,
                                    asset=repo_asset
                                )
                                self.engine.events.put(repo_order)
                                break
                        
                        # 买入债券
                        self.place_order(symbol,
                                        quantity,
                                        OrderType.MARKET, 
                                        None, 
                                        1, 
                                        AssetType.BOND)
                        # order = OrderEvent(
                        #     type=EventType.ORDER,
                        #     timestamp=current_time,
                        #     asset_type=AssetType.BOND,
                        #     symbol=symbol,
                        #     quantity=quantity,
                        #     direction=1,  # 1表示买入
                        #     order_type=OrderType.MARKET,
                        #     asset=asset
                        # )
                        # self.engine.events.put(order)
                        available_cash -= quantity * current_price_full
                        # logger.info(
                        #     f"下单买入债券: {symbol}, "
                        #     f"最大可买金额={max_position:.2f}, "
                        #     f"买入数量={quantity}, "
                        #     f"价格={current_price:.4f}, "
                        #     f"买入后剩余现金={available_cash:.2f}"
                        # )
                
                elif signal < 0 and current_quantity > 0:  # 卖出信号且有持仓
                    # 卖出全部持仓
                    order = OrderData(
                        type=EventType.ORDER,
                        timestamp=current_time,
                        asset_type=AssetType.BOND,
                        symbol=symbol,
                        quantity=current_quantity,
                        direction=-1,  # -1表示卖出
                        order_type=OrderType.MARKET,
                        asset=asset
                    )
                    self.engine.events.put(order)
                    logger.info(
                        f"下单卖出债券: {symbol}, "
                        f"卖出数量={current_quantity}, "
                        f"价格={current_price:.4f}"
                    )
    
    def on_cashflow(self, event: CashFlowEvent) -> None:
        """
        处理现金流事件
        
        Args:
            event: 现金流事件
        """
        if event.cashflow_type == CashFlowType.PRINCIPAL and \
            event.asset_type == AssetType.REPO:
            if event.amount <= 0:  # 到期资金流出，则表示到期标的为正回购
                direction = -1
            else:  # 到期资金流入，则表示到期标的为负回购
                direction = 1
            # 如果回购到期，执行回购融资
            quantity = abs(event.amount)
            repo_order = OrderData(
                type=EventType.ORDER,
                timestamp=event.timestamp,
                asset_type=AssetType.REPO,
                symbol='repo_2',
                quantity=quantity,
                direction=direction,
                order_type=OrderType.MARKET,
                asset=self.engine.assets.get('repo_2')
            )
            self.engine.events.put(repo_order)
